from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
import base64
import random
import string
'''
RSA:
    使用最广泛的非对称加密算法
    公钥（Public Key）:用于加密
    私钥（Private Key）: 保密，用于解密
    秘钥长度： 128,256,384,512 字节
    能够加密的明文长度有限，短于秘钥长度
    同样的秘钥和明文，每次加密得到不一样的密文
'''
# # use random module generate a random string.
# def generate_random_string(l=20):
#     charater1 = string.ascii_letters + string.digits
#     res_s = ''.join([random.choice(charater1) for i in range(l)])
#     return res_s
# # a = generate_random_string(50)
# # print(a)


# generate key
key = RSA.generate(2048)  # 指定秘钥位的长度，8位（bit） = 1字节（byte）,所以会得到一个长度为256字节的秘钥
# save Private key
with open('private_key.pem','wb') as f:
    f.write(key.export_key(format='PEM'))
print('private key already save as private_key.pem.')
# save Public key
with open('public_key.pem','wb') as f:
    f.write(key.publickey().export_key(format='PEM'))
print('public key already save as public_key.pem.')

# use the key 
# define a encrypt string.
string1 = '好好学习asdf123'
with open('public_key.pem','r') as file:
    public_key = RSA.import_key(file.read())
en_obj = PKCS1_OAEP.new(public_key)
en_res = base64.b64encode(en_obj.encrypt(string1.encode('utf-8')))  # encode将字符转为2进制类型数据，最后转为base64编码
en_string= en_res.decode()
print(en_string)
# 解密
with open('private_key.pem','r') as file:
    private_key = RSA.import_key(file.read())
assert public_key.n == private_key.n, "密钥不匹配！"
en_obj = PKCS1_OAEP.new(private_key)
encrypted_bytes = base64.b64decode(en_string.encode())
res = en_obj.decrypt(encrypted_bytes)
print(res.decode())